import argparse
import json
import os,utils
import tqdm
from datetime import datetime
import logging
import asyncio
import argparse
import json
import os
import asyncio
import aiohttp
import tqdm.asyncio as tqdm_async
from typing import Dict, Any


# 禁用所有不必要的日志
logging.getLogger("openai").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("httpcore").setLevel(logging.WARNING)
logging.getLogger("urllib3").setLevel(logging.WARNING)


global project_name 


# 并发数
semaphore = asyncio.Semaphore(30)  # 最大并发数

example = '''example 1:
    question:
    static int ossl_statem_server13_read_transition(SSL_CONNECTION *s, int mt)
{
    OSSL_STATEM *st = &s->statem;

    /*
     * Note: There is no case for TLS_ST_BEFORE because at that stage we have
     * not negotiated TLSv1.3 yet, so that case is handled by
     * ossl_statem_server_read_transition()
     */

    switch (st->hand_state) {
    default:
        break;

    case TLS_ST_EARLY_DATA:
        if (s->hello_retry_request == SSL_HRR_PENDING) {
            if (mt == SSL3_MT_CLIENT_HELLO) {
                st->hand_state = TLS_ST_SR_CLNT_HELLO;
                return 1;
            }

    case TLS_ST_CR_CERT_REQ:
        if (s->post_handshake_auth == SSL_PHA_REQUESTED) {
            if (do_compressed_cert(s))
                st->hand_state = TLS_ST_CW_COMP_CERT;
            else
                st->hand_state = TLS_ST_CW_CERT;
            return WRITE_TRAN_CONTINUE;
        
    Reasoning process:
    This is a server read transition function (ossl_statem_server13_read_transition), so it processes incoming messages from clients. The function checks message type (mt) parameter. In the first case, it checks for SSL3_MT_CLIENT_HELLO, which indicates a CLIENT_HELLO message being received. In the second case, it's preparing to send certificate messages, so the direction is SEND. The function name contains "read_transition" which indicates RECEIVE direction for incoming messages.

    answer:

{
    "function": "ossl_statem_server13_read_transition",
    "transitions": [
        {
            "start_state": "TLS_ST_EARLY_DATA",
            "condition": "s->hello_retry_request == SSL_HRR_PENDING and mt == SSL3_MT_CLIENT_HELLO",
            "condition explanation": "server need client send Hello Retry Request __AND__ server read message is SSL3_MT_CLIENT_HELLO",
            "message_type": "CLIENT_HELLO",
            "message_direction": "RECEIVE",
            "end_state": "TLS_ST_SR_CLNT_HELLO"
        },
        {
            "start_state": "TLS_ST_CR_CERT_REQ",
            "condition": "s->post_handshake_auth == SSL_PHA_REQUESTED and do_compressed_cert(s)",
            "condition explanation":"post handshake authentication was requested __AND__ compressed certificate was requested",
            "message_type": "CERTIFICATE",
            "message_direction": "SEND",
            "end_state": "TLS_ST_CW_COMP_CERT"
        },
        {
            "start_state": "TLS_ST_CR_CERT_REQ",
            "condition": "s->post_handshake_auth == SSL_PHA_REQUESTED and not do_compressed_cert(s)",
            "condition explanation":"post handshake authentication was requested __AND__ __NOT__ compressed certificate was requested",
            "message_type": "CERTIFICATE",
            "message_direction": "SEND",
            "end_state": "TLS_ST_CW_CERT"
        }
    ]
}


            '''


def get_fsm_prompt(protocol_impl, func_data):
    """生成FSM提取提示词"""
    return f'''You are an experienced protocol expert. The protocol implementation is {protocol_impl}. Extract state transitions from the function code with message information.

For each transition, you need to intelligently identify:
1. **Message Type**: Analyze the function name, code context, and any message type constants (like SSL3_MT_*, TLS_MT_*, etc.) to determine the TLS message type. Common types include: CLIENT_HELLO, SERVER_HELLO, ENCRYPTED_EXTENSIONS, CERTIFICATE, CERTIFICATE_VERIFY, FINISHED, KEY_EXCHANGE, HELLO_DONE, CHANGE_CIPHER_SPEC, HANDSHAKE_COMPLETE, etc.

2. **Message Direction**: Determine if the message is being SENT or RECEIVED based on:
   - Function name patterns (e.g., "read_transition" = RECEIVE, "write_transition" = SEND, "Do*" = RECEIVE, "Send*" = SEND)
   - Code context (checking message type constants vs setting state for sending)
   - State transition patterns (preparing to send vs processing received message)

The output format should include: transitions (array with start_state, condition, condition explanation, message_type, message_direction, end_state). The logical connectives in condition explanation should be replaced with: {{"and":"__AND__", "or":"__OR__","not":"__NOT__"}}. Think step by step and analyze the code carefully. {example}

question: {func_data}'''  




def setup_logger(log_path):
    """Set up logging configuration"""
    # 创建 logger 对象（而不是直接用 basicConfig）
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)  # 全局最低级别设为 DEBUG
    
    # 清除默认的 handlers（避免重复日志）
    logger.handlers.clear()
    
    # 1. 文件处理器（记录 DEBUG 及以上级别）
    file_handler = logging.FileHandler(log_path, encoding="utf-8")
    file_handler.setLevel(logging.DEBUG)  # 文件记录所有 DEBUG+ 日志
    file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(file_formatter)
    
    # 2. 控制台处理器（仅记录 INFO 及以上级别）
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)  # 终端只打印 INFO+
    console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(console_formatter)
    
    # 添加处理器
    logger.addHandler(file_handler)
    # logger.addHandler(console_handler)
    
    return logger
            


async def process_function(group_name: str, func_name: str, func_data: Dict, 
                          protocol_impl: str, logger):
    """处理单个函数的异步任务"""

    async with semaphore:  # 控制并发数
        try:
            # 将函数信息传递给LLM，让LLM自动判断消息类型和方向
            enhanced_func_data = {
                "function_name": func_name,
                "code": func_data
            }
            
            prompt = get_fsm_prompt(protocol_impl, enhanced_func_data)
            logger.info(f"Sending to LLM: {func_name}")
            logger.info(f"{func_name}: {prompt}")

            answer = await utils.get_ans_json_async(prompt)
            
            logger.debug(f"Response for {func_name}:\n{json.dumps(answer, indent=2)}")
            return (group_name, func_name, answer)
            
        except Exception as e:
            logger.error(f"Error in {func_name}: {str(e)}")
            return (group_name, func_name, {"error": str(e)})


def get_function(config, func_name):
    """
    从配置指定的项目中获取特定函数的源代码
    
    Args:
        config: 包含项目配置的字典
        func_name: 要提取的函数名称
    
    Returns:
        str: 函数的源代码内容，如果找不到则返回None
    """
    try:
        # 获取项目目录和函数信息
        project_dir = config.get("project_dir")
        state_functions = config.get("state_functions", {})
        
        # 在所有状态函数中查找目标函数
        target_func_info = None
        for state_name, group_data in state_functions.items():
            if func_name in group_data:
                target_func_info = group_data[func_name]
                break
        
        if not target_func_info:
            print(f"函数 {func_name} 未在配置中找到")
            return None
        
        # 构建完整文件路径
        relative_path = target_func_info["path"]
        full_path = os.path.join(project_dir, relative_path)
        
        # 检查文件是否存在
        if not os.path.exists(full_path):
            print(f"文件不存在: {full_path}")
            return None
        
        # 读取指定行号范围的内容
        start_line = target_func_info["start"]
        end_line = target_func_info["end"]
        
        return extract_code_snippet(full_path, start_line, end_line)
        
    except Exception as e:
        print(f"获取函数时出错: {str(e)}")
        return None

# 移除硬编码的消息类型提取函数，让LLM自动判断


def concise_function(state_var, function_code):
    """
    检查函数代码中是否有对state_var的赋值操作，
    并返回包含这些赋值的代码片段（智能扩展切片范围，最多100行）。
    
    参数:
        state_var: 要检查的变量名（如'ssl->options.serverState'）
        function_code: 函数的源代码字符串
        
    返回:
        生成器，每次生成一个包含赋值操作的代码片段
    """
    if function_code is None:
        return
    
    lines = function_code.split('\n')
    var_parts = state_var.replace('->', '.').split('.')
    var_name = var_parts[-1]  # 提取变量名最后部分
    
    # 构建更精确的匹配模式
    # 对于 s->statem.hand_state，我们寻找 st->hand_state = 或 st.hand_state = 这样的赋值
    if len(var_parts) >= 2:
        # 构建精确的匹配模式：var_part[-2].var_part[-1] = 或 var_part[-2]->var_part[-1] =
        pattern1 = f"{var_parts[-2]}.{var_parts[-1]}"
        pattern2 = f"{var_parts[-2]}->{var_parts[-1]}"
        # 对于OpenSSL，还需要匹配 st->hand_state 格式
        pattern3 = f"st->{var_parts[-1]}"
    else:
        # 如果只有一个部分，直接使用
        pattern1 = var_parts[-1]
        pattern2 = var_parts[-1]
        pattern3 = f"st->{var_parts[-1]}"
    
    # 找到所有赋值操作的行号
    assignment_lines = []
    for i, line in enumerate(lines):
        stripped = line.strip()
        # 跳过注释行
        if stripped.startswith('#') or stripped.startswith('//'):
            continue
            
        # 检查简单赋值 - 使用更精确的匹配
        if '=' in line:
            left_side = line.split('=')[0].strip()
            # 检查是否包含完整的状态变量赋值模式
            if ((pattern1 in left_side and '=' in line) or 
                (pattern2 in left_side and '=' in line) or 
                (pattern3 in left_side and '=' in line)):
                assignment_lines.append(i)
        
        # 检查复合赋值
        for op in ['+=', '-=']:
            if op in line:
                left_side = line.split(op)[0].strip()
                if ((pattern1 in left_side) or 
                    (pattern2 in left_side) or 
                    (pattern3 in left_side)):
                    assignment_lines.append(i)
                    break
    
    # 合并相邻的赋值区域
    merged_ranges = []
    if assignment_lines:
        current_start = assignment_lines[0]
        current_end = assignment_lines[0]
        
        for line in assignment_lines[1:]:
            if line <= current_end + 50:  # 如果下一个赋值在50行内，扩展范围
                current_end = line
            else:
                merged_ranges.append((current_start, current_end))
                current_start = current_end = line
        merged_ranges.append((current_start, current_end))
    
    # 生成代码片段
    for start, end in merged_ranges:
        # 计算扩展范围，确保不超过100行
        context_size = min(25, 100 // (end - start + 1)) if start != end else 25
        expanded_start = max(0, start - context_size)
        expanded_end = min(len(lines), end + context_size + 1)
        
        # 确保总行数不超过100
        while expanded_end - expanded_start > 100:
            if expanded_end - start > end - expanded_start:
                expanded_end -= 1
            else:
                expanded_start += 1
        
        yield '\n'.join(lines[expanded_start:expanded_end])



def extract_code_snippet(file_path, start_line, end_line):
    """
    从文件中提取指定行号范围的代码
    
    Args:
        file_path: 文件路径
        start_line: 起始行号
        end_line: 结束行号
    
    Returns:
        str: 提取的代码内容
    """
    try:
        extracted_lines = []
        
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
            for current_line_num, line in enumerate(file, 1):
                if current_line_num >= start_line and current_line_num <= end_line:
                    extracted_lines.append(line)
                elif current_line_num > end_line:
                    break
        
        return ''.join(extracted_lines)
        
    except FileNotFoundError:
        print(f"文件未找到: {file_path}")
        return None
    except Exception as e:
        print(f"读取文件时出错: {str(e)}")
        return None






async def async_main(config):
    """异步主函数"""
    global project_name
    save_dir = os.path.join(os.path.abspath(config['save_parent_dir']), 
                          f"{project_name}_{config['version']}")
    os.makedirs(save_dir, exist_ok=True)
    
    log_path = os.path.join(save_dir, 'get_fsm_from_llm.log')
    logger = setup_logger(log_path)
    logger.info(f"Starting async processing for {project_name}")
    
    # 创建aiohttp客户端会话
    async with aiohttp.ClientSession() as session:
        tasks = []
        print(f"开始处理 {len(config['state_functions'])} 个状态组")
        for group_name, group_data in config['state_functions'].items():
            print(f"处理组: {group_name}")
            # 获取状态变量和枚举
            state_var = group_data.get('state_var')
            state_enum = group_data.get('state_enum', [])
            
            # 遍历组中的函数
            for func_name, func_info in group_data.items():
                if func_name in ['state_var', 'state_enum']:
                    continue
                
                function_code = get_function(config, func_name)
                if function_code is None:
                    logger.warning(f"跳过函数 {func_name}，无法获取代码")
                    continue
                
                snippet_count = 0
                for snippet in concise_function(state_var, function_code):
                    snippet_count += 1
                    task = process_function(
                        group_name, func_name, 
                        snippet, config['project'],
                        logger
                    )
                    tasks.append(task)
                
                print(f"  {func_name}: {snippet_count} 个代码片段")
        
        # 使用tqdm显示进度
        results = []
        for f in tqdm_async.tqdm.as_completed(tasks, desc="Processing"):
            try:
                result = await f
                results.append(result)
            except Exception as e:
                logger.error(f"Task failed: {str(e)}")
        
        # 整理结果
        llm_fsm_response = {}
        for group_name, func_name, answer in results:
            if group_name not in llm_fsm_response:
                llm_fsm_response[group_name] = {}
            llm_fsm_response[group_name][func_name] = answer
        
        # 保存结果
        save_path = os.path.join(save_dir, "get_fsm_from_llm.json")
        with open(save_path, 'w', encoding='utf-8') as f:
            json.dump(llm_fsm_response, f, indent=2)
        
        logger.info(f"Saved results to {save_path}")


# python .\async_get_fsm_from_llm2.py .\config\wolfssl.json
if __name__ == "__main__":

        
    # Windows平台特殊处理
    if os.name == 'nt':
        asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())


    parser = argparse.ArgumentParser(description='Sync code analysis')
    parser.add_argument('config', help='Path to config JSON')
    config = parser.parse_args().config

    # config = r".\config\test.json"

    with open(config, encoding='utf-8') as f:
        config = json.load(f)

    project_name = config['project']
    
    asyncio.run(async_main(config))

    # 在文件最后加
    import warnings
    warnings.filterwarnings("ignore", category=RuntimeWarning, message=".*Event loop is closed.*")